# 这部分代码在上面代码的基础上对文本进行分割处理
import pandas as pd
from sklearn.model_selection import train_test_split
import argparse


def split_text(text, label, max_length):
    """
    根据max_length分割text，并返回分割后的列表和对应的label列表。

    参数:
        text (str): 空格分隔的数字字符串
        label (int): 类别标签
        max_length (int): 最大长度

    返回:
        tuple: 包含分割后的文本列表和对应的标签列表
    """
    tokens = text.split()
    token_count = len(tokens)

    if token_count <= max_length:
        return [(text, label)]

    segments = []
    for i in range(0, token_count, max_length):
        segment = tokens[i:i + max_length]
        if len(segment) == max_length:
            segments.append((' '.join(segment), label))
        else:
            segment = tokens[-max_length:]
            segments.append((' '.join(segment), label))
            break

    return segments


def process_file(filename, max_length):
    """
    处理指定的CSV文件，对文本进行分段。

    参数:
        filename (str): 输入文件名
        max_length (int): 最大文本长度
    """
    # 读取CSV文件
    df = pd.read_csv(filename, sep='\t')

    # 创建一个新的DataFrame来存储处理后的数据
    new_rows = []

    for index, row in df.iterrows():
        text = row['text']
        label = row['label']

        # 分割文本并添加到新行列表中
        segments = split_text(text, label, max_length)
        new_rows.extend(segments)

    # 将新行转换为DataFrame
    df = pd.DataFrame(new_rows, columns=['text', 'label'])

    # 保存处理后的数据
    df.to_csv(filename, sep='\t', index=False)
    print(f"处理后的数据已保存到 {filename}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="对文本数据进行分段处理")
    parser.add_argument("--filename", type=str, required=True, help="要处理的CSV文件路径")
    parser.add_argument("--max_length", type=int, default=4095, help="最大文本长度")
    
    args = parser.parse_args()
    
    process_file(args.filename, args.max_length)
